#include <device/blkio.h>

#include <mm/slab.h>
#include <mm/valloc.h>

#include <arch/cpu.h>

#include <compiler.h>
#include <timer.h>

#define LOG_TAG "blkio"
#include <log.h>

static kmem_cache_t* ioreq_kmem_cache = NULL;

void blkio_init() {
    ioreq_kmem_cache = new_kmem_cache("blkio_req", sizeof(blkio_req_t), 1, 0);
}

void timer_io_req_cb(void* data) {
    blkio_ctx_t* ctx = data;
    blkio_schedule(ctx);
}

bool blkio_deadlock_checker(void* data) {
    blkio_req_t* req = container_of(data, blkio_req_t, wait);

    if ((req->flags & BLKIO_BUSY) && list_empty(&req->reqs)) {
        if (req->is_req_done(req->track_data)) {
            blkio_schedule(req->io_ctx);
            // blkio_complete(req);
            return true;
        }
    }

    return false;
}

static inline blkio_req_t* blkio_req_create(pbuffer_t* buffer, u64 start_lba, u32 options) {
    options = options & ~0xF;
    blkio_req_t* breq = take_obj(ioreq_kmem_cache);
    *breq = (blkio_req_t){
        .blk_addr = start_lba,
        .flags = options,
    };
    breq->pb = buffer;
    waitq_init(&breq->wait, blkio_deadlock_checker);
    return breq;
}

blkio_req_t* blkio_read_req(pbuffer_t* pb, u64 start_lba, u32 opt) {
    return blkio_req_create(pb, start_lba, opt);
}

blkio_req_t* blkio_write_req(pbuffer_t* pb, u64 start_lba, u32 opt) {
    blkio_req_t* breq = blkio_req_create(pb, start_lba, opt);
    breq->flags |= BLKIO_WRITE;
    return breq;
}

void blkio_free_req(blkio_req_t* req) {
    release_obj(ioreq_kmem_cache, req);
}

blkio_ctx_t* blkio_newctx(io_req_handler handler) {
    blkio_ctx_t* ctx = vzalloc(sizeof(blkio_ctx_t));
    ctx->req_handle = handler;

    list_head_init(&ctx->queue);
    mutex_init(&ctx->lock);

    timer_create(0, TIMER_PERIODIC, timer_io_req_cb, ctx);

    return ctx;
}

void blkio_commit(blkio_ctx_t* ctx, blkio_req_t* req, int options) {
    if (req->flags & BLKIO_PENDING) {
        return;
    }

    req->flags |= BLKIO_PENDING;
    req->io_ctx = ctx;

    if ((options & BLKIO_WAIT)) {
        prepare_to_wait(&req->wait);
    }

    mutex_lock(&ctx->lock);

    list_append(&ctx->queue, &req->reqs);

    mutex_unlock(&ctx->lock);

    if (!ctx->busy) {
        if ((options & BLKIO_WAIT)) {
            blkio_schedule(ctx);
            try_wait(&req->wait);
            return;
        }
        blkio_schedule(ctx);
    } else if ((options & BLKIO_WAIT)) {
        try_wait(&req->wait);
    }
}

void blkio_schedule(blkio_ctx_t* ctx) {
    if (mutex_locked(&ctx->lock)) {
        return;
    }

    mutex_lock(&ctx->lock);
    if (list_empty(&ctx->queue)) {
        mutex_unlock(&ctx->lock);
        return;
    }

    blkio_req_t* head = container_of(ctx->queue.next, blkio_req_t, reqs);
    list_remove(&head->reqs);

    head->flags |= BLKIO_BUSY;
    head->flags &= ~BLKIO_PENDING;
    head->io_ctx->busy++;

    mutex_unlock(&ctx->lock);

    ctx->req_handle(head);
}

void blkio_complete(blkio_req_t* req) {
    req->flags &= ~(BLKIO_BUSY | BLKIO_PENDING);

    wake_all(&req->wait);

    if ((req->flags & BLKIO_FREE)) {
        blkio_free_req(req);
    }

    if ((req->errcode)) {
        LOGE("blkio complete with error: %d", req->errcode);
    }

    req->io_ctx->busy--;
}